import os

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, TensorDataset
from torch.nn.utils.rnn import pad_sequence
# 创建滑动窗口的数据集
# acce: ax, ay, az, gx, gy, gz (N, 6) angle: (N,) 对齐时间戳
def data_align(imu_data, gps_data):
    imu_time = imu_data[:,0]
    gps_time = gps_data[:,0]
    imu_slices = []

    for i in range(len(gps_time)-1):
        t1 = gps_time[i]
        t2 = gps_time[i+1]

        idex_t1 = np.searchsorted(imu_time, t1, side='left')
        idex_t2 = np.searchsorted(imu_time, t2, side='right')

        idx_min = max(0, idex_t1 - 50)
        idx_max = min(len(imu_time), idex_t2 + 50)

        temp_slice = imu_data[idx_min:idx_max]
        imu_slices.append(temp_slice)
    return imu_slices, gps_data

def read_data_file(head):
    loca_path = os.path.join(head, 'relative_polar_coordinates.csv')
    acce_path = os.path.join(head, 'Accelerometer.csv')
    gyra_path = os.path.join(head, 'Gyroscope.csv')
    # 时间戳一致，直接读一个就行
    loca = pd.read_csv(loca_path)[['time','distance','angle']].to_numpy()
    acce = pd.read_csv(acce_path)[['time','x','y','z']].to_numpy()
    gyra = pd.read_csv(gyra_path)[['x','y','z']].to_numpy()

    # time, ax, ay, az, gx, gy, gz
    imu_data = np.concatenate([acce, gyra], axis=1)
    gps_data = loca
    '''
    # check
    print(imu_data.shape)
    print(gps_data.shape)
    print(imu_data[0])
    print(gps_data[0])'''
    return imu_data, gps_data

def get_dataloader(batch_size):

    imu_tensor_list = []
    gps_tensor_list = []
    dir = '../dataset'
    data = os.listdir(dir)
    for d in data:
        head = os.path.join(dir,d)

        #head = './dataset/No_81_Denggao_Road-2024-10-03_05-19-03'
        imu_data, gps_data = read_data_file(head)
        imu_slices, gps_data = data_align(imu_data, gps_data)
        #assert len(imu_slices) == len(gps_data)
        imu_slices_tensor = [torch.tensor(slice, dtype=torch.float32) for slice in imu_slices]
        padded_imu_slices = pad_sequence(imu_slices_tensor, batch_first=True)[:, :, 1:]
        gps_data = torch.tensor(gps_data, dtype=torch.float32)[1:,2]
        imu_tensor_list.append(padded_imu_slices)
        gps_tensor_list.append(gps_data)
    max_m = max(t.shape[1] for t in imu_tensor_list)
    padded_tensors = []
    for t in imu_tensor_list:
        current_m = t.shape[1]
        pad_size = max_m - current_m
        padded_t = torch.nn.functional.pad(t, (0, 0, 0, pad_size))
        padded_tensors.append(padded_t)
    padded_imu_slices = torch.cat(padded_tensors, dim=0)
    gps_data = torch.cat(gps_tensor_list, dim=0)
    #print(padded_imu_slices.shape)

    #print(gps_data.shape)

    divi = int(0.1 * padded_imu_slices.shape[0])
    test_set = TensorDataset(padded_imu_slices[:divi,::], gps_data[:divi])
    train_set = TensorDataset(padded_imu_slices[divi:,::], gps_data[divi:])
    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=True)
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
    return train_loader, test_loader


#dataloader = get_dataloader(32)
